package io.jjz;

import com.cmb.pb.Example;
import com.google.protobuf.RpcCallback;
import com.google.protobuf.RpcController;
import com.google.protobuf.Service;
import org.apache.hadoop.hbase.Cell;
import org.apache.hadoop.hbase.CellUtil;
import org.apache.hadoop.hbase.Coprocessor;
import org.apache.hadoop.hbase.CoprocessorEnvironment;
import org.apache.hadoop.hbase.client.Scan;
import org.apache.hadoop.hbase.coprocessor.CoprocessorException;
import org.apache.hadoop.hbase.coprocessor.CoprocessorService;
import org.apache.hadoop.hbase.coprocessor.RegionCoprocessorEnvironment;
import org.apache.hadoop.hbase.filter.FirstKeyOnlyFilter;
import org.apache.hadoop.hbase.protobuf.ResponseConverter;
import org.apache.hadoop.hbase.regionserver.InternalScanner;
import org.apache.hadoop.hbase.regionserver.RegionScanner;
import org.apache.hadoop.hbase.util.Bytes;

import java.io.IOException;
import java.util.ArrayList;

/**
 * Author       : 蒋俊钊 80256973
 * Date         : 2019/2/21
 * Copyright    (C) ChinaMerchantsBank
 */
public class RowCountEndpoint extends Example.RowCountService implements Coprocessor, CoprocessorService {

    private RegionCoprocessorEnvironment env;

    @Override
    public void start(CoprocessorEnvironment env) throws IOException {
        if (env instanceof RegionCoprocessorEnvironment) {
            this.env = (RegionCoprocessorEnvironment) env;
        } else {
            throw new CoprocessorException("Must be loader on a table region!");
        }
    }

    @Override
    public void stop(CoprocessorEnvironment env) throws IOException {

    }

    @Override
    public Service getService() {
        return this;
    }

    @Override
    public void getRowCount(RpcController controller, Example.CountRequest request, RpcCallback<Example.CountResponse> done) {
        Scan scan = new Scan();
        scan.setFilter(new FirstKeyOnlyFilter());
        Example.CountResponse response = null;
        InternalScanner scanner = null;

        try {
            scanner = env.getRegion().getScanner(scan);
            ArrayList<Cell> results = new ArrayList<>();
            boolean hasMore;
            byte[] lastRow = null;
            long count = 0;
            do {
                hasMore = scanner.next(results);
                for (Cell c : results) {
                    byte[] currentRow = CellUtil.cloneRow(c);
                    if (lastRow == null || !Bytes.equals(lastRow, currentRow)) {
                        lastRow = currentRow;
                        count++;
                    }
                }
                results.clear();
            } while (hasMore);
            response = Example.CountResponse.newBuilder().setCount(count).build();

        } catch (IOException e) {
            ResponseConverter.setControllerException(controller, e);
        } finally {
            if (scanner != null) {
                try {
                    scanner.close();
                } catch (IOException e) {
                }
            }
        }
        done.run(response);
    }
}
